package com.zzx.schoolfellow.config;

import com.zzx.schoolfellow.model.entity.ChatMessage;
import org.springframework.web.socket.CloseStatus;
import org.springframework.web.socket.TextMessage;
import org.springframework.web.socket.WebSocketSession;
import org.springframework.web.socket.handler.TextWebSocketHandler;

import java.util.List;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.CopyOnWriteArrayList;

/**
 * @program: schoolfeoolw_backend
 * @description:
 * @author: zzx
 * @create: 2023-03-26 12:45
 **/
public class WebSocketHandler extends TextWebSocketHandler {

    private final Map<String, WebSocketSession> sessions = new ConcurrentHashMap<>();

    @Override
    public void afterConnectionEstablished(WebSocketSession session) throws Exception {
        String userId = getUserId(session);
        sessions.put(userId, session);
    }

    @Override
    protected void handleTextMessage(WebSocketSession session, TextMessage message) throws Exception {
        String userId = getUserId(session);
        ChatMessage chatMessage = ChatMessage.fromJson(message.getPayload());
        String targetUserId = chatMessage.getReceiver();
        WebSocketSession targetSession = sessions.get(targetUserId);
        if (targetSession != null && targetSession.isOpen()) {
            targetSession.sendMessage(new TextMessage(chatMessage.toJson()));
        }
    }

    @Override
    public void afterConnectionClosed(WebSocketSession session, CloseStatus status) throws Exception {
        String userId = getUserId(session);
        sessions.remove(userId);
    }

    private String getUserId(WebSocketSession session) {
        Map<String, Object> attributes = session.getAttributes();
        String userId = (String) attributes.get("userId");
        if (userId == null) {
            throw new IllegalStateException("WebSocket session中没有设置用户ID");
        }
        return userId;
    }


}
